from langgraph.prebuilt import create_react_agent
from langgraph_swarm import create_swarm, create_handoff_tool
from major.models_manager import qwen_model

# 预订酒店
def book_hotel(hotel_name: str):
    """预订酒店"""
    print("假装成功预定了酒店")
    return f"成功预订了 {hotel_name} 的住宿。"

# 预订航班
def book_flight(from_airport: str, to_airport: str):
    """预订航班"""
    print("假装成功预定了航班")
    return f"成功预订了从 {from_airport} 到 {to_airport} 的航班。"

# 创建工具，用于代理之间的交接
transfer_to_hotel_assistant = create_handoff_tool(
    agent_name="hotel_assistant"
)
transfer_to_flight_assistant = create_handoff_tool(
    agent_name="flight_assistant"
)

# 创建航班助手代理
flight_assistant = create_react_agent(
    model=qwen_model.get_model(),
    tools=[book_flight, transfer_to_hotel_assistant],
    prompt="你是一个航班预订助手。",
    name="flight_assistant"
)

# 创建酒店助手代理
hotel_assistant = create_react_agent(
    model=qwen_model.get_model(),
    tools=[book_hotel, transfer_to_flight_assistant],
    prompt="你是一个酒店预订助手。",
    name="hotel_assistant"
)

# 创建 swarm 代理，协调多个代理的工作
swarm = create_swarm(
    model=qwen_model.get_model(),
    agents=[flight_assistant, hotel_assistant],
    default_active_agent="flight_assistant"
).compile()

for chunk in swarm.stream(
    {
        "messages": [
            {
                "role": "user",
                "content": "预订一个从 BOS 到 JFK 的航班，并预订 McKittrick Hotel 的住宿。"
            }
        ]
    }
):
    print(chunk)
    print("\n")

